import tkinter as tk
from tkinter import messagebox, filedialog
import threading
import os
import json
import time
from typing import Dict, Any

from models.factory import ModelFactory
from ui.main_window import MainWindow
from utils.config_manager import ConfigManager
from utils.chat_history import ChatHistoryManager
from utils.knowledge_base import KnowledgeBase

# 新增导入
from ui.discussion_tab import DiscussionTab

class ChatApp:
    """聊天应用程序主类"""
    
    def __init__(self):
        self.root = tk.Tk()
        
        # 初始化各个管理器
        self.model_factory = ModelFactory()
        self.config_manager = ConfigManager("config.json")
        self.chat_history = ChatHistoryManager("chat_history.json")
        self.knowledge_base = KnowledgeBase(dir_path="knowledge_base")
        
        # 获取模型配置
        self.models = self.model_factory.get_model_configs()
        
        # 从配置中加载API密钥
        for model_name in self.models:
            api_key = self.config_manager.get_api_key(model_name)
            if api_key:
                self.models[model_name]["api_key"] = api_key
        
        # 初始化UI
        self.window = MainWindow(self.root, self.config_manager)
        
        # 设置回调函数
        self.window.set_callbacks(
            on_send=self.handle_send_message,
            on_model_change=self.handle_model_change,
            on_thinking_change=self.handle_thinking_change,
            on_save_settings=self.handle_save_settings,
            on_clear_chat=self.handle_clear_chat,
            on_role_change=self.handle_role_change,
            on_export_history=self.handle_export_history,
            on_import_history=self.handle_import_history,
            on_clear_history=self.handle_clear_history
        )
        
        # 设置讨论回调函数
        if hasattr(self.window, 'chat_tab'):
            self.window.chat_tab.on_start_discussion = self.handle_start_discussion
            self.window.chat_tab.on_view_discussion_progress = self.handle_view_discussion_progress
        
        # 当前选中的模型
        self.current_model = list(self.models.keys())[0]
        self.window.get_model_var().set(self.current_model)
        
        # 深度思考设置
        self.enable_thinking = False
        
        # 加载配置
        self.load_config()
        
        # 新增属性
        self.typing_speed = self.config_manager.get_typing_speed()
        self.is_typing = False
        self.typing_thread = None
        
        # 讨论相关属性
        self.discussion_tab = self.window.discussion_tab
        self.discussion_history = []
    
    def handle_send_message(self, message: str) -> None:
        """处理发送消息事件"""
        # 显示用户消息
        self.window.add_message(message, "user")
        self.chat_history.add_message("user", message)
        
        # 在新线程中获取AI响应
        threading.Thread(target=self.get_ai_response, args=(message,), daemon=True).start()
    
    def handle_model_change(self) -> None:
        """处理模型切换事件"""
        model_var = self.window.get_model_var()
        if model_var:
            new_model = model_var.get()
            if new_model and new_model in self.models:
                self.current_model = new_model
                # 保存当前选中的模型
                self.config_manager.save_current_model(new_model)
        
        # 更新API密钥输入框
        # self.window.api_key_entry.delete(0, tk.END)
        # self.window.api_key_entry.insert(0, self.models[self.current_model]["api_key"])
        
        # 更新深度思考开关状态
        if not self.models[self.current_model]["supports_thinking"]:
            self.window.get_thinking_var().set(False)
            self.enable_thinking = False
    
    def handle_thinking_change(self) -> None:
        """处理深度思考开关事件"""
        if not self.models[self.current_model]["supports_thinking"] and self.window.get_thinking_var().get():
            self.window.get_thinking_var().set(False)
            return
        self.enable_thinking = self.window.get_thinking_var().get()
    
    def handle_save_settings(self) -> None:
        """处理保存设置事件"""
        self.save_all_settings()
    
    def handle_clear_chat(self) -> None:
        """处理清空聊天事件"""
        # 清空当前会话
        self.chat_history.clear_current_session()
        
        # 删除聊天历史文件
        history_file_path = os.path.join(os.path.dirname(__file__), "chat_history.json")
        if os.path.exists(history_file_path):
            try:
                os.remove(history_file_path)
            except OSError as e:
                print(f"删除聊天历史文件时出错: {e}")
        
        # 清空界面显示
        self.window.clear_chat()
    
    def load_config(self) -> None:
        """加载配置"""
        # 准备设置数据
        settings = {
            'role': self.config_manager.get_setting("role", "助手"),
            'max_tokens': self.config_manager.get_setting("max_tokens", 2000),
            'temperature': self.config_manager.get_setting("temperature", 0.7),
            'system_prompt': self.config_manager.get_setting("system_prompt", None),
            'typing_speed': self.config_manager.get_typing_speed()
        }
        
        # 加载设置到窗口
        self.window.load_settings(settings)
        
        # 加载当前选中的模型
        current_model = self.config_manager.get_current_model()
        if current_model and current_model in self.models:
            self.current_model = current_model
            self.window.get_model_var().set(self.current_model)
            # 触发模型初始化
            self.window.initialize_chat_model()
        
        # 加载当前选中的角色
        current_role = self.config_manager.get_current_role()
        if hasattr(self.window, 'get_role_var'):
            role_var = self.window.get_role_var()
            if role_var and current_role:
                role_var.set(current_role)
    
    def save_all_settings(self) -> None:
        """保存所有设置"""
        try:
            # 获取设置
            settings = self.window.get_settings()
            
            # 保存API密钥
            for model_name in self.models:
                if f'{model_name}_api_key' in settings:
                    api_key = settings[f'{model_name}_api_key']
                    self.models[model_name]["api_key"] = api_key
                    self.config_manager.save_api_key(model_name, api_key)
            
            # 保存其他设置
            if 'role' in settings:
                self.config_manager.save_setting("role", settings['role'])
            if 'max_tokens' in settings:
                self.config_manager.save_setting("max_tokens", settings['max_tokens'])
            if 'temperature' in settings:
                self.config_manager.save_setting("temperature", settings['temperature'])
            if 'system_prompt' in settings:
                self.config_manager.save_setting("system_prompt", settings['system_prompt'])
            if 'typing_speed' in settings:
                self.config_manager.save_typing_speed(settings['typing_speed'])
                self.typing_speed = settings['typing_speed']
            
            messagebox.showinfo("成功", "设置已保存！")
        except Exception as e:
            messagebox.showerror("错误", f"保存设置时出错: {str(e)}")
    
    def handle_role_change(self) -> None:
        """处理角色切换事件"""
        # 从聊天页面获取角色变更
        if hasattr(self.window, 'get_role_var'):
            role_var = self.window.get_role_var()
            if role_var:
                current_role = role_var.get()
                if current_role:
                    # 保存当前角色
                    self.config_manager.save_current_role(current_role)
                    
                    # 获取角色配置并保存系统提示词
                    role_config = self.config_manager.get_role_config(current_role)
                    if 'system_prompt' in role_config:
                        self.config_manager.save_setting("system_prompt", role_config['system_prompt'])
                        # 同时更新聊天页面的current_role_system_prompt属性
                        if hasattr(self.window, 'chat_tab'):
                            self.window.chat_tab.current_role_system_prompt = role_config['system_prompt']
                    
                    return
        
        # 从设置页面获取角色变更（保持原有逻辑）
        settings = self.window.get_settings()
        if 'role' in settings:
            self.config_manager.save_setting("role", settings['role'])
            
            # 同时保存更新后的系统
            
            
            if 'system_prompt' in settings:
                self.config_manager.save_setting("system_prompt", settings['system_prompt'])
                # 同时更新聊天页面的current_role_system_prompt属性
                if hasattr(self.window, 'chat_tab'):
                    self.window.chat_tab.current_role_system_prompt = settings['system_prompt']
    
    def handle_export_history(self) -> None:
        """处理导出历史事件"""
        try:
            file_path = filedialog.asksaveasfilename(
                defaultextension=".json",
                filetypes=[("JSON files", "*.json"), ("All files", "*.*")],
                title="导出聊天历史"
            )
            if file_path:
                history_data = self.chat_history.get_all_history()
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(history_data, f, ensure_ascii=False, indent=2)
                messagebox.showinfo("成功", f"聊天历史已导出到: {file_path}")
        except Exception as e:
            messagebox.showerror("错误", f"导出历史时出错: {str(e)}")
    
    def handle_import_history(self) -> None:
        """处理导入历史事件"""
        try:
            file_path = filedialog.askopenfilename(
                filetypes=[("JSON files", "*.json"), ("All files", "*.*")],
                title="导入聊天历史"
            )
            if file_path:
                with open(file_path, 'r', encoding='utf-8') as f:
                    history_data = json.load(f)
                self.chat_history.load_from_data(history_data)
                self.update_history_display()
                messagebox.showinfo("成功", "聊天历史已导入！")
        except Exception as e:
            messagebox.showerror("错误", f"导入历史时出错: {str(e)}")
    
    def handle_clear_history(self) -> None:
        """处理清空历史事件"""
        result = messagebox.askyesno("确认", "确定要清空所有聊天历史吗？此操作不可撤销。")
        if result:
            self.chat_history.clear_all_history()
            self.update_history_display()
            messagebox.showinfo("成功", "聊天历史已清空！")
    
    def update_history_display(self) -> None:
        """更新历史显示"""
        history = self.chat_history.get_all_history()
        self.window.set_history_data(history)
        self.window.update_history_display()
    
    def get_ai_response(self, user_message: str) -> None:
        """获取AI响应"""
        try:
            # 先尝试从知识库匹配答案
            match_result = self.knowledge_base.match_question(user_message)
            if match_result:
                matched_q, answers = match_result
                # 从答案列表中随机选择一个
                import random
                response = random.choice(answers)
                self.root.after(0, lambda: self.window.add_message(response, "assistant"))
                self.chat_history.add_message("assistant", response)
                return
            
            # 知识库没有匹配结果，使用AI模型回答
            # 检查API密钥
            current_api_key = self.config_manager.get_api_key(self.current_model)
            if not current_api_key or len(current_api_key.strip()) < 10:
                response = f"请先为模型 {self.current_model} 配置有效的API密钥。"
                self.root.after(0, lambda: self.window.add_message(response, "system"))
                self.chat_history.add_message("system", response)
                return
            
            # 创建模型实例
            model = self.model_factory.create_model(self.models[self.current_model])
            
            # 构建消息
            # 优先使用聊天页面中设置的角色系统提示词
            system_prompt = None
            if hasattr(self.window, 'chat_tab') and hasattr(self.window.chat_tab, 'current_role_system_prompt'):
                system_prompt = self.window.chat_tab.current_role_system_prompt
            
            # 如果聊天页面没有设置，则从配置管理器获取
            if not system_prompt:
                system_prompt = self.config_manager.get_setting("system_prompt", "你是一个友善的AI助手，请用中文回答问题。")
                role = self.config_manager.get_setting("role", "助手")
                if role != "助手":
                    system_prompt = f"你是一个{role}，{system_prompt}"
            
            if self.enable_thinking and model.supports_thinking:
                system_prompt += "请在回答前进行深度思考，展示你的思考过程。"
            
            # 添加系统消息和历史消息
            messages = [{"role": "system", "content": system_prompt}]
            # 过滤掉历史消息中的system和assistant消息，避免重复和隐私问题
            history_context = [msg for msg in self.chat_history.get_context() if msg["role"] not in ["system", "assistant"]]
            messages.extend(history_context)
            messages.append({"role": "user", "content": user_message})
            
            # 获取响应
            response = model.generate_response(messages, self.enable_thinking)
            
            # 处理响应
            if isinstance(response, dict) and "thinking" in response:
                # 有思考过程的响应
                thinking_process = response["thinking"]
                result = response["response"]
                self.chat_history.add_message("assistant", result)
                self.root.after(0, lambda: self.window.add_message(
                     result, "assistant", thinking_content=thinking_process
                 ))
            else:
                # 普通响应
                self.chat_history.add_message("assistant", response)
                self.root.after(0, lambda: self.window.add_message(
                     response, "assistant"
                 ))
            
            # 保存聊天历史
            self.chat_history.save_history()
            
        except Exception as e:
            error_msg = f"获取AI响应时出错: {str(e)}"
            self.root.after(0, lambda: self.window.add_message("system", error_msg))
    
    def run(self) -> None:
        """运行应用程序"""
        # 显示欢迎消息
        welcome_msg = "欢迎使用AI聊天助手！请在设置中配置您的API密钥，然后开始聊天。"
        self.root.after(1000, lambda: self.window.add_message(welcome_msg, "system"))
        self.chat_history.add_message("system", welcome_msg)
        
        # 更新历史显示
        self.root.after(1500, self.update_history_display)
        
        # 启动主循环
        self.root.mainloop()

    def handle_start_discussion(self, topic: str, config: Dict[str, Any]) -> None:
        """处理开始讨论事件"""
        # 在新线程中执行讨论
        threading.Thread(target=self._run_discussion, args=(topic, config), daemon=True).start()
    
    def handle_view_discussion_progress(self) -> None:
        """处理查看讨论进度事件"""
        # 切换到讨论标签页
        if self.discussion_tab:
            self.window.notebook.select(self.discussion_tab.parent_frame)
    
    def _run_discussion(self, topic: str, config: Dict[str, Any]) -> None:
        """执行AI讨论"""
        try:
            # 获取配置
            rounds = config.get('rounds', 3)
            scheduler_model = config.get('scheduler_model')
            participant_models = config.get('participant_models', {})
            
            # 创建讨论记录
            discussion_id = int(time.time())
            discussion_record = {
                'id': discussion_id,
                'topic': topic,
                'date': time.strftime('%Y-%m-%d %H:%M:%S'),
                'rounds': rounds,
                'scheduler_model': scheduler_model,
                'participant_models': participant_models,
                'messages': []
            }
            
            # 添加讨论记录
            self.discussion_history.append(discussion_record)
            
            # 创建调度模型实例
            scheduler_config = self.models[scheduler_model].copy()
            scheduler_config['api_key'] = self.config_manager.get_api_key(scheduler_model)
            scheduler = self.model_factory.create_model(scheduler_config)
            
            # 创建参与者模型实例
            participants = {}
            for model_name, count in participant_models.items():
                if count > 0:
                    model_config = self.models[model_name].copy()
                    model_config['api_key'] = self.config_manager.get_api_key(model_name)
                    participants[model_name] = {
                        'config': model_config,
                        'model': self.model_factory.create_model(model_config),
                        'count': count
                    }
            
            # 执行多轮讨论
            for round_num in range(rounds):
                # 调度模型生成提示
                prompt = self._generate_scheduler_prompt(topic, round_num, rounds, participants)
                
                # 添加调度消息
                discussion_record['messages'].append({
                    'round': round_num + 1,
                    'role': 'scheduler',
                    'model': scheduler_model,
                    'content': prompt
                })
                
                # 更新讨论显示
                self._update_discussion_display(discussion_id)
                
                # 参与者模型响应
                for model_name, participant_info in participants.items():
                    for i in range(participant_info['count']):
                        response = self._get_participant_response(
                            participant_info['model'], prompt, topic, round_num
                        )
                        
                        # 添加参与者消息
                        discussion_record['messages'].append({
                            'round': round_num + 1,
                            'role': 'participant',
                            'model': model_name,
                            'instance': i + 1,
                            'content': response
                        })
                        
                        # 更新讨论显示
                        self._update_discussion_display(discussion_id)
                        
                        # 增加延迟，避免触发API频率限制
                        time.sleep(3)
                
                # 增加每轮之间的延迟
                time.sleep(5)
            
            # 最终总结
            summary = self._generate_final_summary(scheduler, topic, discussion_record['messages'])
            discussion_record['summary'] = summary
            
            # 更新讨论显示
            self._update_discussion_display(discussion_id)
            
            # 显示完成消息
            self.root.after(0, lambda: self.window.add_message(
                f"🤖 AI讨论已完成！点击'查看讨论进度'查看详细结果", "system"
            ))
            
        except Exception as e:
            error_msg = f"讨论过程中出错: {str(e)}"
            self.root.after(0, lambda: self.window.add_message(error_msg, "system"))
    
    def _generate_scheduler_prompt(self, topic: str, round_num: int, total_rounds: int, participants: Dict) -> str:
        """生成调度提示"""
        participant_list = []
        for model_name, info in participants.items():
            participant_list.append(f"{model_name} ({info['count']}个实例)")
        
        # 根据轮次生成不同的提示词
        if round_num == 0:
            # 第一轮：问题分析和初步答案
            prompt = f"""
第 {round_num + 1}/{total_rounds} 轮讨论 - 问题分析
讨论主题: {topic}
参与者: {', '.join(participant_list)}

请直接针对讨论主题进行分析，并给出该问题的答案。
请重点关注：
1. 问题的核心要点和关键因素
2. 可能的解决方向和初步答案
3. 需要进一步探讨的具体问题

请提供具体、实用的分析和给出该问题的答案：
"""
        elif round_num == total_rounds - 1:
            # 最后一轮：总结和最终答案
            prompt = f"""
第 {round_num + 1}/{total_rounds} 轮讨论 - 最终答案
讨论主题: {topic}
参与者: {', '.join(participant_list)}

基于前几轮的讨论，请给出该问题的答案和具体实施方案。
请提供：
1. 完整的答案框架
2. 具体的实施步骤
3. 预期的效果和注意事项
4. 后续跟进答案

请提供详细、可操作的最终答案：
"""
        else:
            # 中间轮次：深入讨论和答案优化
            prompt = f"""
第 {round_num + 1}/{total_rounds} 轮讨论 - 深入探讨
讨论主题: {topic}
参与者: {', '.join(participant_list)}

请在前一轮讨论的基础上，进一步深入探讨并优化答案。
请重点关注：
1. 答案的可行性和实施细节
2. 潜在的问题和应对策略
3. 不同角度的优化答案
4. 具体的改进措施

请提供深入的优化答案和具体改进：
"""
        
        return prompt.strip()
    
    def _get_participant_response(self, model, prompt: str, topic: str, round_num: int) -> str:
        """获取参与者响应"""
        messages = [
            {"role": "system", "content": f"你正在参与关于'{topic}'的AI讨论。请尽力解答讨论主题，然后查看别的AI的意见，进行反驳或修改自己的意见。"},
            {"role": "user", "content": prompt}
        ]
        
        response = model.generate_response(messages, False)
        return response
    
    def _generate_final_summary(self, scheduler, topic: str, messages: list) -> str:
        """生成最终总结"""
        # 提取所有参与者的讨论内容
        discussion_content = "\n".join([
            f"{msg['role']} ({msg.get('model', '')}): {msg['content']}"
            for msg in messages
        ])
        
        summary_prompt = f"""
讨论主题: {topic}

以下是完整的讨论记录:
{discussion_content}

请基于以上讨论内容，生成一个全面的答案总结报告。
总结应该专注于提供具体的、可操作的答案，包括：

1. 最终答案框架 - 完整的解决方法和策略
2. 具体实施步骤 - 详细的、分步骤的执行计划
3. 资源需求和时间安排 - 所需资源和时间估算
4. 预期效果和衡量指标 - 可量化的成果预期
5. 风险分析和应对措施 - 潜在问题及答案
6. 后续跟进答案 - 持续改进和监控机制

请确保总结是实用、可执行的，避免理论性描述，专注于具体行动方案：
"""
        
        messages = [
            {"role": "system", "content": "你是一个专业的讨论总结者，请基于讨论记录生成高质量的解决方案总结。"},
            {"role": "user", "content": summary_prompt}
        ]
        
        summary = scheduler.generate_response(messages, False)
        return summary
    
    def _update_discussion_display(self, discussion_id: int) -> None:
        """更新讨论显示"""
        # 找到对应的讨论记录
        discussion = None
        for d in self.discussion_history:
            if d['id'] == discussion_id:
                discussion = d
                break
        
        if discussion and self.discussion_tab:
            # 检查是否已经添加过这个讨论
            existing_discussion = None
            for d in self.discussion_tab.discussions:
                if d['id'] == discussion_id:
                    existing_discussion = d
                    break
            
            if existing_discussion:
                # 更新现有的讨论记录
                self.root.after(0, lambda: self.discussion_tab.update_discussion(discussion))
            else:
                # 添加新的讨论记录
                self.root.after(0, lambda: self.discussion_tab.add_discussion(discussion))

def main():
    """主函数"""
    app = ChatApp()
    app.run()

if __name__ == "__main__":
    main()